package com.buddy.summary.code.tree;

import org.junit.Test;

public class AvlBinarySearchTree<AnyType extends Comparable<? super AnyType>> {

    private static class AvlNode<AnyType> {
        AnyType element;
        AvlNode left;
        AvlNode right;
        int height;

        public AvlNode(AnyType theElement) {
            this(theElement, null, null);
        }

        public AvlNode(AnyType theElement, AvlNode<AnyType> lt, AvlNode<AnyType> rt) {
            element = theElement;
            left = lt;
            right = rt;
            height = 0;
        }

        public AvlNode(AnyType element, AvlNode left, AvlNode right, int height) {
            this.element = element;
            this.left = left;
            this.right = right;
            this.height = height;
        }
    }

    private AvlNode<AnyType> root;

    public void insert(AnyType x) {
        root = insert(x, root);
    }

    private AvlNode<AnyType> insert(AnyType x, AvlNode<AnyType> node) {
        if (node == null)
            return new AvlNode<>(x);
        int compareResult = x.compareTo(node.element);
        if (compareResult < 0)
            node.left = insert(x, node.left);
        else if (compareResult > 0)
            node.right = insert(x, node.right);
        return balance(node);
    }

    private int height(AvlNode node) {
        return node == null ? -1 : node.height;
    }

    private static final int ALLOWED_IMBALANCE = 1;

    private AvlNode<AnyType> balance(AvlNode<AnyType> node) {
        if (node == null)
            return node;
        if (height(node.left) - height(node.right) > ALLOWED_IMBALANCE) {
            if (height(node.left.left) >= height(node.left.right))
                node = rotateWithLeftChild(node);
            else
                node = doubleRotateWithLeftChild(node);
        } else if (height(node.right) - height(node.left) > ALLOWED_IMBALANCE) {
            if (height(node.right.right) >= height(node.right.left))
                node = rotateWithRightChild(node);
            else
                node = doubleRotateWithRightChild(node);
        }
        node.height = Math.max(height(node.left), height(node.right)) + 1;
        return node;
    }

    /**
     * 左单旋转
     * 子树 左节点变根节点，根节点变右子树
     *
     * @param k2
     * @return
     */
    private AvlNode<AnyType> rotateWithLeftChild(AvlNode<AnyType> k2) {
        AvlNode<AnyType> k1 = k2.left;
        k2.left = k1.right;
        k1.right = k2;
        k2.height = Math.max(height(k2.left), height(k2.right)) + 1;
        k1.height = Math.max(height(k1.left), height(k1.right)) + 1;
        return k1;
    }

    /**
     * 右右单旋转
     * 子树右节点变根节点，根节点变左子树
     * @param k2
     * @return
     */
    private AvlNode<AnyType> rotateWithRightChild(AvlNode<AnyType> k2) {
        AvlNode<AnyType> k1 = k2.right;
        k2.right = k1.left;
        k1.left = k2;
        k2.height = Math.max(height(k2.left), height(k2.right)) + 1;
        k1.height = Math.max(height(k1.left), height(k1.right)) + 1;
        return k1;
    }

    /**
     * 左右双旋转
     * @param k3
     * @return
     */
    private AvlNode<AnyType> doubleRotateWithLeftChild(AvlNode<AnyType> k3) {
        k3.left = rotateWithRightChild(k3.left);
        return rotateWithLeftChild(k3);
    }
    /**
     * 右左双旋转
     * @param k3
     * @return
     */
    private AvlNode<AnyType> doubleRotateWithRightChild(AvlNode<AnyType> k3) {
        k3.right = rotateWithLeftChild(k3.left);
        return rotateWithRightChild(k3);
    }
    @Test
    public void avlBinarySearchTreeTest() {
        AvlBinarySearchTree<Integer> tree = new AvlBinarySearchTree<>();
        for (int i = 3; i >= 1; i--) {
            tree.insert(i);
            System.out.println(tree.root.element);
        }
        for (int i = 4; i <= 7; i++) {
            tree.insert(i);
            System.out.println(tree.root.element);
        }
    }
}
